import akshare as ak
import sqlite3
import pandas as pd
import os
from datetime import datetime
from datacache import table_exists

conn = sqlite3.connect(os.path.join(os.getcwd(), 'stock_data_cache.db'), check_same_thread=False)

# 获取中证500或者沪深300的数据
def get_all_index_stock_data(index_code):
    start_date_base = '20100101'
    end_date_base = datetime.now().strftime('%Y%m%d')
    df = ak.index_zh_a_hist(symbol=index_code, period='daily', start_date=start_date_base, end_date=end_date_base)
    # 如果 DataFrame 不为空，则将其保存到 SQLite 数据库
    if not df.empty:
        table_name = f'index_stocks_{index_code}'
        # 目标表已存在，则替换
        df.to_sql(table_name, conn, if_exists='replace', index=False)
        print(f"Data for {index_code} cached in table '{table_name}'.")

# 获取中证500或者沪深300的数据  # 沪深300为 000300 中证500 000905
def get_index_stock_data(index_code, start_date, end_date):
    index_table_name = f'index_stocks_{index_code}'
    is_table_exist = table_exists(index_table_name)
    if not is_table_exist:
        get_all_index_stock_data(index_code)
    try:
        # 连接到 SQLite 数据库
        query = f"""
        SELECT *
        FROM { index_table_name }
        WHERE 日期 BETWEEN ? AND ?
        ORDER BY 日期 ASC
        """
        # 使用 pandas 读取 SQL 查询结果
        df = pd.read_sql_query(query, conn, params=(start_date, end_date))
        return df
    except sqlite3.OperationalError:
        print(f"No cached data found for symbol {index_code}.")
        return None

# 根据所属行业查询股票
def get_stock_by_type(stock_type):
    query = f"SELECT * FROM sz_stocks WHERE 所属行业 = ?"
    df = pd.read_sql_query(query, conn, params=(stock_type,))
    return df

def get_sz_stocks(plate):
    is_table_exist = table_exists("sz_stocks")
    if not is_table_exist:
        stock_info_sh_df = ak.stock_info_sz_name_code()
        stock_info_sh_df.to_sql("sz_stocks", conn, if_exists='replace', index=False)
    query = f"SELECT * FROM sz_stocks where 板块 = ?"
    df = pd.read_sql_query(query, conn, params=(plate,))
    return df

def get_sh_stocks():
    is_table_exist = table_exists("sh_stocks")
    if not is_table_exist:
        stock_info_sh_df = ak.stock_info_sh_name_code()
        print(stock_info_sh_df)
        stock_info_sh_df.to_sql("sh_stocks", conn, if_exists='replace', index=False)
    query = f"SELECT * FROM sh_stocks"
    df = pd.read_sql_query(query, conn)
    return df


if __name__ == "__main__":
    get_sz_stocks("")